import geopandas as gpd
import pandas as pd
import numpy as np
import os
import gc
import math
from shapely.geometry import Polygon, MultiPolygon
from shapely.validation import make_valid
import re
from datetime import datetime

def export_fengshan_yulin_model(geojson_path, output_obj_path, scale=100000, thickness=2.5):
    """
    封山育林模型导出工具：
    1. 适配"封山育林"GeoJSON的字段结构（LIN_CHANG/LIN_BAN/XIAO_BAN）
    2. 优化几何处理，解决模型破碎问题
    3. 生成带有厚度的3D模型，确保相邻地块无缝贴合
    """
    print(f"🌲 开始处理封山育林数据...")
    print(f"📁 输入文件: {geojson_path}")
    print(f"📁 输出文件: {output_obj_path}")
    print(f"📏 模型厚度: {thickness} 单位 | 缩放比例: {scale}")
    print(f"⏰ 开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

    # 检查输入文件是否存在
    if not os.path.exists(geojson_path):
        print(f"❌ 错误：输入文件不存在 - {geojson_path}")
        return False

    try:
        # 1. 读取并预处理数据
        gdf = gpd.read_file(geojson_path)
        print(f"✅ 成功读取数据：共 {len(gdf)} 个地块")

        # 验证核心字段
        required_fields = ["LIN_CHANG", "LIN_BAN", "XIAO_BAN", "geometry"]
        for field in required_fields:
            if field not in gdf.columns:
                print(f"⚠️  警告：缺少必要字段 '{field}'，将使用默认值")

        # 数据清洗与标准化
        valid_plots = []
        for idx, row in gdf.iterrows():
            # 提取并格式化核心属性（适配封山育林数据字段）
            林场 = str(row.get("LIN_CHANG", "未知林场")).strip()
            林班 = str(row.get("LIN_BAN", "0000")).strip().zfill(4)
            小班 = str(row.get("XIAO_BAN", "000")).strip().zfill(3)
            面积 = float(row.get("MIAN_JI", 0))  # 提取面积信息

            # 清理特殊字符（确保文件名合法）
            林场 = re.sub(r'[\\/*?:"<>|]', '', 林场).replace(' ', '_')
            完整名称 = f"{林场}_林班{林班}_小班{小班}"

            # 处理几何数据
            几何 = row.geometry
            if 几何 is None or 几何.is_empty:
                print(f"⚠️  跳过空几何：{完整名称}")
                continue

            # 修复无效几何
            if not 几何.is_valid:
                几何 = make_valid(几何)
                print(f"🔧 修复无效几何：{完整名称}")

            # 过滤过小地块（面积小于0.1公顷的忽略）
            if 面积 > 0.1 or (面积 == 0 and 几何.area > 1e-7):
                valid_plots.append({
                    "林场": 林场,
                    "林班": 林班,
                    "小班": 小班,
                    "完整名称": 完整名称,
                    "面积": 面积,
                    "几何": 几何
                })
            else:
                print(f"⚠️  跳过过小地块（面积 {面积} 公顷）：{完整名称}")

        if not valid_plots:
            print(f"❌ 没有有效地块数据可处理")
            return False
        print(f"✅ 筛选后有效地块：{len(valid_plots)} 个")

        # 2. 建立层级结构（林场→林班→小班）
        hierarchy = {}
        for plot in valid_plots:
            林场 = plot["林场"]
            林班 = plot["林班"]

            if 林场 not in hierarchy:
                hierarchy[林场] = {}
            if 林班 not in hierarchy[林场]:
                hierarchy[林场][林班] = []

            hierarchy[林场][林班].append(plot)

        print(f"📊 层级结构：{len(hierarchy)} 个林场 → {sum(len(v) for v in hierarchy.values())} 个林班")

        # 3. 收集所有坐标点用于全局处理
        all_coords = []
        for plot in valid_plots:
            geom = plot["几何"]
            processed_geoms = process_geometry(geom)
            for g in processed_geoms:
                # 提取经纬度坐标（忽略Z值）
                exterior_coords = [(p[0], p[1]) for p in g.exterior.coords]
                all_coords.extend(exterior_coords)
                for interior in g.interiors:
                    interior_coords = [(p[0], p[1]) for p in interior.coords]
                    all_coords.extend(interior_coords)

        # 计算全局偏移（减少坐标误差）
        coords_array = np.array(all_coords, dtype=np.float64)
        min_lon, min_lat = np.min(coords_array, axis=0)
        max_lon, max_lat = np.max(coords_array, axis=0)
        global_offset = np.array([(min_lon + max_lon) / 2, (min_lat + max_lat) / 2])
        print(f"🌍 地理范围：\n   经度：{min_lon:.6f} ~ {max_lon:.6f}\n   纬度：{min_lat:.6f} ~ {max_lat:.6f}")

        # 创建坐标映射表（确保相邻地块共享顶点）
        coord_map = create_coord_map(all_coords, precision=8)
        print(f"🔗 坐标映射表：{len(coord_map)} 个唯一坐标点")

        # 4. 生成OBJ数据
        obj_data = []
        total_vertices = 0
        total_faces = 0
        removed_faces = 0

        # 文件头信息
        obj_data.append(f"# 封山育林3D模型")
        obj_data.append(f"# 生成时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        obj_data.append(f"# 来源：{os.path.basename(geojson_path)}")
        obj_data.append(f"# 包含：{len(hierarchy)} 林场 | {len(valid_plots)} 地块")
        obj_data.append("")

        # 写入顶点数据（顶部z=0，底部z=-thickness）
        vertices = []
        sorted_coords = sorted(coord_map.items(), key=lambda x: x[1])
        for (lon, lat), idx in sorted_coords:
            # 经纬度转平面坐标（应用偏移和缩放）
            x = (lon - global_offset[0]) * scale
            y = (lat - global_offset[1]) * scale

            # 顶部顶点
            obj_data.append(f"v {x:.8f} {y:.8f} 0.0")
            vertices.append((x, y, 0.0))

            # 底部顶点
            obj_data.append(f"v {x:.8f} {y:.8f} {-thickness}")
            vertices.append((x, y, -thickness))

        total_vertices = len(vertices)
        obj_data.append(f"\n# 顶点总数：{total_vertices}")
        obj_data.append("")

        # 生成面数据
        print(f"🔨 正在生成面数据...")
        for 林场名, 林班_dict in hierarchy.items():
            obj_data.append(f"o {林场名}")
            obj_data.append(f"usemtl {林场名}_material")
            obj_data.append(f"# 林场：{林场名}")
            obj_data.append("")

            for 林班名, 小班列表 in 林班_dict.items():
                林班_fullname = f"{林场名}_林班{林班名}"
                obj_data.append(f"g {林班_fullname}")
                obj_data.append(f"usemtl {林班_fullname}_material")
                obj_data.append(f"# 林班：{林班名}（{len(小班列表)}个小班）")
                obj_data.append("")

                for 小班 in 小班列表:
                    小班名 = 小班["完整名称"]
                    obj_data.append(f"g {小班名}")
                    obj_data.append(f"usemtl {小班名}_material")
                    obj_data.append(f"# 小班：{小班名}（面积：{小班['面积']}公顷）")

                    # 处理几何并生成面
                    geom = 小班["几何"]
                    processed_geoms = process_geometry(geom)
                    for g_idx, g in enumerate(processed_geoms):
                        faces, removed = generate_optimized_faces(
                            g, coord_map, vertices,
                            min_area=0.1, precision=8
                        )
                        removed_faces += removed

                        if not faces:
                            print(f"⚠️  {小班名}的子几何{g_idx}无有效面，已跳过")
                            continue

                        # 写入面数据
                        for face in faces:
                            obj_data.append(f"f {face[0]} {face[1]} {face[2]}")
                        total_faces += len(faces)

                    obj_data.append("")  # 小班间分隔
                obj_data.append("")  # 林班间分隔
            obj_data.append("")  # 林场间分隔

        # 5. 写入OBJ文件
        os.makedirs(os.path.dirname(output_obj_path), exist_ok=True)
        with open(output_obj_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(obj_data))

        # 6. 输出统计信息
        print(f"\n🎉 模型导出完成！")
        print(f"📊 统计信息：")
        print(f"   - 顶点数：{total_vertices:,}")
        print(f"   - 有效面数：{total_faces:,}")
        print(f"   - 移除无效面：{removed_faces}")
        print(f"   - 输出文件：{output_obj_path}")
        print(f"⏰ 完成时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

        gc.collect()
        return True

    except Exception as e:
        print(f"❌ 导出失败：{str(e)}")
        import traceback
        traceback.print_exc()
        return False

def create_coord_map(coords, precision=8):
    """创建坐标映射表，确保相同坐标使用同一顶点索引"""
    coord_map = {}
    idx = 0
    for (lon, lat) in coords:
        # 四舍五入到指定精度，消除浮点误差
        key = (round(lon, precision), round(lat, precision))
        if key not in coord_map:
            coord_map[key] = idx
            idx += 1
    return coord_map

def process_geometry(geom):
    """处理几何对象，确保多边形闭合并拆分多面体"""
    processed = []
    if isinstance(geom, Polygon):
        if not geom.is_empty:
            # 确保多边形闭合
            if not np.allclose(geom.exterior.coords[0][:2], geom.exterior.coords[-1][:2], atol=1e-8):
                exterior = list(geom.exterior.coords) + [geom.exterior.coords[0]]
                interiors = [list(interior.coords) for interior in geom.interiors]
                processed.append(Polygon(exterior, interiors))
            else:
                processed.append(geom)
    elif isinstance(geom, MultiPolygon):
        # 拆分多面体为单个多边形
        for poly in geom.geoms:
            if not poly.is_empty:
                if not np.allclose(poly.exterior.coords[0][:2], poly.exterior.coords[-1][:2], atol=1e-8):
                    exterior = list(poly.exterior.coords) + [poly.exterior.coords[0]]
                    interiors = [list(interior.coords) for interior in poly.interiors]
                    processed.append(Polygon(exterior, interiors))
                else:
                    processed.append(poly)
    return processed

def generate_optimized_faces(geom, coord_map, vertices, min_area=0.1, precision=8):
    """生成优化的面数据，过滤无效面和重复面"""
    faces = []
    removed = 0
    face_set = set()  # 用于检测重复面

    # 处理外部轮廓
    exterior_coords = [(p[0], p[1]) for p in geom.exterior.coords]
    top_indices = []
    bottom_indices = []

    for (lon, lat) in exterior_coords:
        key = (round(lon, precision), round(lat, precision))
        if key not in coord_map:
            continue

        coord_idx = coord_map[key]
        top_indices.append(coord_idx * 2 + 1)    # 顶部顶点索引（OBJ索引从1开始）
        bottom_indices.append(coord_idx * 2 + 2) # 底部顶点索引

    # 生成顶部面
    if len(top_indices) >= 3:
        # 顶部三角化
        for i in range(1, len(top_indices) - 1):
            face = (top_indices[0], top_indices[i], top_indices[i+1])
            face_sorted = tuple(sorted(face))

            if face_sorted in face_set:
                removed += 1
                continue

            # 计算面面积，过滤过小的面
            area = calculate_face_area(face, vertices)
            if area < min_area:
                removed += 1
                continue

            faces.append(face)
            face_set.add(face_sorted)

        # 生成底部面（反向三角化避免背面剔除）
        for i in range(len(bottom_indices) - 2, 0, -1):
            face = (bottom_indices[0], bottom_indices[i], bottom_indices[i+1])
            face_sorted = tuple(sorted(face))

            if face_sorted in face_set:
                removed += 1
                continue

            area = calculate_face_area(face, vertices)
            if area < min_area:
                removed += 1
                continue

            faces.append(face)
            face_set.add(face_sorted)

        # 生成侧面（连接顶部和底部）
        for i in range(len(top_indices) - 1):
            v1 = top_indices[i]
            v2 = top_indices[i+1]
            v3 = bottom_indices[i+1]
            v4 = bottom_indices[i]

            # 第一个三角形
            face1 = (v1, v2, v3)
            face1_sorted = tuple(sorted(face1))
            if face1_sorted not in face_set and calculate_face_area(face1, vertices) >= min_area:
                faces.append(face1)
                face_set.add(face1_sorted)
            else:
                removed += 1

            # 第二个三角形
            face2 = (v1, v3, v4)
            face2_sorted = tuple(sorted(face2))
            if face2_sorted not in face_set and calculate_face_area(face2, vertices) >= min_area:
                faces.append(face2)
                face_set.add(face2_sorted)
            else:
                removed += 1

    # 处理内部孔洞
    for interior in geom.interiors:
        interior_coords = [(p[0], p[1]) for p in interior.coords]
        hole_top = []
        hole_bottom = []

        for (lon, lat) in interior_coords:
            key = (round(lon, precision), round(lat, precision))
            if key not in coord_map:
                continue

            coord_idx = coord_map[key]
            hole_top.append(coord_idx * 2 + 1)
            hole_bottom.append(coord_idx * 2 + 2)

        if len(hole_top) >= 3:
            # 孔洞顶部面（反向三角化）
            for i in range(len(hole_top) - 2, 0, -1):
                face = (hole_top[0], hole_top[i], hole_top[i+1])
                face_sorted = tuple(sorted(face))
                if face_sorted in face_set or calculate_face_area(face, vertices) < min_area:
                    removed += 1
                    continue
                faces.append(face)
                face_set.add(face_sorted)

            # 孔洞底部面（正向三角化）
            for i in range(1, len(hole_bottom) - 1):
                face = (hole_bottom[0], hole_bottom[i], hole_bottom[i+1])
                face_sorted = tuple(sorted(face))
                if face_sorted in face_set or calculate_face_area(face, vertices) < min_area:
                    removed += 1
                    continue
                faces.append(face)
                face_set.add(face_sorted)

            # 孔洞侧面
            for i in range(len(hole_top) - 1):
                v1 = hole_top[i]
                v2 = hole_top[i+1]
                v3 = hole_bottom[i+1]
                v4 = hole_bottom[i]

                face1 = (v2, v1, v3)
                face1_sorted = tuple(sorted(face1))
                if face1_sorted not in face_set and calculate_face_area(face1, vertices) >= min_area:
                    faces.append(face1)
                    face_set.add(face1_sorted)
                else:
                    removed += 1

                face2 = (v3, v1, v4)
                face2_sorted = tuple(sorted(face2))
                if face2_sorted not in face_set and calculate_face_area(face2, vertices) >= min_area:
                    faces.append(face2)
                    face_set.add(face2_sorted)
                else:
                    removed += 1

    return faces, removed

def calculate_face_area(face, vertices):
    """计算三角形面的面积，用于过滤过小的面"""
    try:
        # 顶点索引转换（OBJ索引从1开始）
        v1 = vertices[face[0] - 1]
        v2 = vertices[face[1] - 1]
        v3 = vertices[face[2] - 1]

        # 提取X、Y坐标
        x1, y1, _ = v1
        x2, y2, _ = v2
        x3, y3, _ = v3

        # 使用 shoelace 公式计算面积
        return 0.5 * abs((x2 - x1) * (y3 - y1) - (y2 - y1) * (x3 - x1))
    except:
        return 0  # 错误时返回0，确保被过滤

# 执行入口
if __name__ == "__main__":
    # 配置文件路径
    input_geojson = r"C:\Users\ma\Desktop\鹤北geojson\封山育林.geojson"  # 替换为你的输入文件路径
    output_obj = r"C:\Users\ma\Desktop\鹤北geojson\封山育林_3D模型.obj"    # 替换为你的输出文件路径

    # 执行导出
    success = export_fengshan_yulin_model(
        geojson_path=input_geojson,
        output_obj_path=output_obj,
        scale=100000,    # 坐标缩放比例
        thickness=3.0    # 模型厚度，可根据需要调整
    )

    if success:
        print("\n🎉 封山育林3D模型导出成功！")
    else:
        print("\n❌ 导出失败，请检查错误信息")
